{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the 30 minutes Flower federated learning tutorial!\n",
    "\n",
    "In this tutorial you will implement your first Federated Learning project using [Flower](https://flower.ai/).\n",
    "\n",
    "🧑‍🏫 This tutorial starts at zero and expects no familiarity with federated learning. Only a basic understanding of data science and Python programming is assumed. A minimal understanding of ML is not required but if you already know about it, nothing is stopping your from modifying this code as you see fit!\n",
    "\n",
    "> [Star Flower on GitHub](https://github.com/adap/flower) ⭐️ and join the Flower community on Flower Discuss and the Flower Slack to connect, ask questions, and get help:\n",
    "> - [Join Flower Discuss](https://discuss.flower.ai/) We'd love to hear from you in the `Introduction` topic! If anything is unclear, post in `Flower Help - Beginners`.\n",
    "> - [Join Flower Slack](https://flower.ai/join-slack) We'd love to hear from you in the `#introductions` channel! If anything is unclear, head over to the `#questions` channel.\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complementary Content\n",
    "\n",
    "But before do so, let me point you to a few video tutorials in the [Flower Youtube channel](https://www.youtube.com/@flowerlabs) that you might want to check out after this tutorial. We post new videos fairly regularly with new content:\n",
    "* **[VIDEO]** quickstart-tensorflow: [15-min video on how to start with Flower + Tensorflow/Keras](https://www.youtube.com/watch?v=FGTc2TQq7VM)\n",
    "* **[VIDEO]** quickstart-pytorch: [20-min video on how to start with Flower + PyTorch](https://www.youtube.com/watch?v=jOmmuzMIQ4c)\n",
    "* **[VIDEO]**  Flower simulation mini-series: [9 line-by-line video tutorials](https://www.youtube.com/watch?v=cRebUIGB5RU&list=PLNG4feLHqCWlnj8a_E1A_n5zr2-8pafTB)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Setup\n",
    "\n",
    "Now let's really begin with this tutorial!\n",
    "\n",
    "To start working with Flower, very little is required once you have activated your Python environment (e.g. via `conda`, `virtualenv`, `pyenv`, etc). If you are running this code on Colab, there is really nothing to do except to install Flower and other dependencies. The steps below have been verified to run in Colab. Let's first, install Flower, then the ML framework of your choice and extra dependencies you might want to use.\n",
    "\n",
    "## Installing Flower\n",
    "\n",
    "You can install flower very conveniently from `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q \"flwr[simulation]\" flwr-datasets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the _simulation_ engine in Flower, which allows you to run a large number of clients without the overheads of manually managing devices. This is achieved via the `Simulation Engine`, the core component in Flower to run simulations efficiently."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install your ML framework\n",
    "\n",
    "Flower is agnostic to your choice of ML Framework. Flower works with `PyTorch`, `Tensorflow`, `NumPy`, `🤗 Transformers`, `MLX`, `JAX`, `scikit-learn`, `fastai`, `Pandas`. Flower also supports all major platforms: `iOS`, `Android` and plain `C++`. You can find a _quickstart- example for each of the above in the [Flower Repository](https://github.com/adap/flower/tree/main/examples) inside the `examples/` directory. And check the [Flower Documentation](https://flower.ai/docs/) for even more learning materials.\n",
    "\n",
    "In this tutorial we are going to use PyTorch, uncomment the line below if you haven't installed PyTorch in your system. In this tutorial we'll use a small model so using CPU only training will suffice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "7192138a-8c87-4d9a-f726-af1038ad264c"
   },
   "outputs": [],
   "source": [
    "# you might see a warning after running the command below, this can be ignored\n",
    "# if you are running this outside Colab, you probably need to adjust the command below\n",
    "# !pip install torch==1.13.1+cpu torchvision==0.14.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to install some other dependencies you are likely familiar with. We'll use these to make plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "58b7af77-609f-4118-bd5b-5629a4b5a296"
   },
   "outputs": [],
   "source": [
    "!pip install matplotlib"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Centralised training: the old way of doing ML"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by creating a simple (but complete) training loop as it is commonly done in centralised setups. Starting our tutorial in this way will allow us to very clearly identify which parts of a typical ML pipeline are common to both centralised and federated training and which ones are poles a part.\n",
    "\n",
    "For this tutorial we'll design a image classification pipeline for [MNIST digits](https://en.wikipedia.org/wiki/MNIST_database) and using a simple CNN model as the network to train. The MNIST dataset is comprised of `28x28` greyscale images with digits from 0 to 9 (i.e. 10 classes in total)\n",
    "\n",
    "\n",
    "## A dataset\n",
    "\n",
    "Let's begin by constructing the dataset. We will use 🤗HuggingFace Datasets to download MNIST. We will prepare a function that will be use later to apply standard normalization transformations from `TorchVision` and create the dataloaders for the `train` and `test` partitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we naturally first need to import torch and torchvision\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.transforms import ToTensor, Normalize, Compose\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "def get_mnist_dataloaders(mnist_dataset, batch_size: int):\n",
    "    pytorch_transforms = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])\n",
    "\n",
    "    # Prepare transformation functions\n",
    "    def apply_transforms(batch):\n",
    "        batch[\"image\"] = [pytorch_transforms(img) for img in batch[\"image\"]]\n",
    "        return batch\n",
    "\n",
    "    mnist_train = mnist_dataset[\"train\"].with_transform(apply_transforms)\n",
    "    mnist_test = mnist_dataset[\"test\"].with_transform(apply_transforms)\n",
    "\n",
    "    # Construct PyTorch dataloaders\n",
    "    trainloader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True)\n",
    "    testloader = DataLoader(mnist_test, batch_size=batch_size)\n",
    "    return trainloader, testloader\n",
    "\n",
    "\n",
    "# Download dataset\n",
    "mnist = load_dataset(\"ylecun/mnist\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can have a quick overview of our datasets by just typing the object on the command line. For instance, below you can see the sizes of both the `train` and `test` partitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "f10b649f-3cee-4e86-c7ff-94bd1fd3e082"
   },
   "outputs": [],
   "source": [
    "mnist"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a more insightful visualisation. First let's see the distribution over the labels by constructing a histogram. Then, let's visualise some training examples !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 490
    },
    "outputId": "c8d0f4c0-60cd-4c58-bc91-3b061dae8046"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "# construct histogram\n",
    "all_labels = mnist[\"train\"][\"label\"]\n",
    "all_label_counts = Counter(all_labels)\n",
    "\n",
    "# visualise histogram\n",
    "bar = plt.bar(all_label_counts.keys(), all_label_counts.values())\n",
    "_ = plt.bar_label(bar)\n",
    "\n",
    "# plot formatting\n",
    "_ = plt.xticks([label for label in all_label_counts.keys()])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualise 32 images from the dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import io\n",
    "\n",
    "\n",
    "def visualise_n_random_examples(trainset_, n: int, verbose: bool = True):\n",
    "    trainset_data = [\n",
    "        Image.open(io.BytesIO(entry[0].as_py())) for entry in trainset_.data[0]\n",
    "    ]\n",
    "    idx = list(range(len(trainset_data)))\n",
    "    random.shuffle(idx)\n",
    "    idx = idx[:n]\n",
    "    if verbose:\n",
    "        print(f\"will display images with idx: {idx}\")\n",
    "\n",
    "    # construct canvas\n",
    "    num_cols = 8\n",
    "    num_rows = int(np.ceil(len(idx) / num_cols))\n",
    "    fig, axs = plt.subplots(figsize=(16, num_rows * 2), nrows=num_rows, ncols=num_cols)\n",
    "\n",
    "    # display images on canvas\n",
    "    for c_i, i in enumerate(idx):\n",
    "        axs.flat[c_i].imshow(trainset_data[i], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 715
    },
    "outputId": "4e0988a8-388d-4acf-882b-089e4ea887bf"
   },
   "outputs": [],
   "source": [
    "# it is likely that the plot this function will generate looks familiar to other plots you might have generated before\n",
    "# or you might have encountered in other tutorials. So far, we aren't doing anything new, Federated Learning will start soon!\n",
    "visualise_n_random_examples(mnist[\"train\"], n=32)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A CNN architecture\n",
    "\n",
    "This tutorial is not so much about novel architectural designs so we keep things simple and make use of a typical CNN that is adequate for the MNIST image classification task.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, num_classes: int) -> None:\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 4 * 4, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, num_classes)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 4 * 4)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly to what we did with the dataset you could inspect the model in various ways. We can, for instance, count the number of model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "67d01ab4-cdd9-4661-8f01-eaa9aabf786d"
   },
   "outputs": [],
   "source": [
    "model = Net(num_classes=10)\n",
    "num_parameters = sum(value.numel() for value in model.state_dict().values())\n",
    "print(f\"{num_parameters = }\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Training Loop\n",
    "\n",
    "A minimal training loop in PyTorch can be constructed with three functions:\n",
    "*  `train()` that will train the model given a dataloader.\n",
    "* `test()` that will be used to evaluate the performance of the model on held-out data, e.g., a training set.\n",
    "* `run_centralised()` which will define additional elements (e.g. the optimiser) and run the training loop over N epochs.\n",
    "\n",
    "Let's construct these functions!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(net, trainloader, optimizer, epochs):\n",
    "    \"\"\"Train the network on the training set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    net.train()\n",
    "    for batch in trainloader:\n",
    "        images, labels = batch[\"image\"], batch[\"label\"]\n",
    "        optimizer.zero_grad()\n",
    "        loss = criterion(net(images), labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "def test(net, testloader):\n",
    "    \"\"\"Validate the network on the entire test set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    correct, loss = 0, 0.0\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in testloader:\n",
    "            images, labels = batch[\"image\"], batch[\"label\"]\n",
    "            outputs = net(images)\n",
    "            loss += criterion(outputs, labels).item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    accuracy = correct / len(testloader.dataset)\n",
    "    return loss, accuracy\n",
    "\n",
    "\n",
    "def run_centralised(\n",
    "    trainloader, testloader, epochs: int, lr: float, momentum: float = 0.9\n",
    "):\n",
    "    \"\"\"A minimal (but complete) training loop\"\"\"\n",
    "\n",
    "    # instantiate the model\n",
    "    model = Net(num_classes=10)\n",
    "\n",
    "    # define optimiser with hyperparameters supplied\n",
    "    optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)\n",
    "\n",
    "    # train for the specified number of epochs\n",
    "    for e in range(epochs):\n",
    "        print(f\"Training epoch {e} ...\")\n",
    "        train(model, trainloader, optim, epochs)\n",
    "\n",
    "    # training is completed, then evaluate model on the test set\n",
    "    loss, accuracy = test(model, testloader)\n",
    "    print(f\"{loss = }\")\n",
    "    print(f\"{accuracy = }\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run this for 5 epochs (you'll see it reaching close to 99% accuracy -- as expected from a centralised setup with the MNIST dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "e8d9b429-178d-4924-e82f-4d4e52863788"
   },
   "outputs": [],
   "source": [
    "# Construct dataloaders\n",
    "trainloader, testloader = get_mnist_dataloaders(mnist, batch_size=32)\n",
    "\n",
    "# Run the centralised training\n",
    "run_centralised(trainloader, testloader, epochs=3, lr=0.01)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above centralised formulation has worked just fine for some applications and to showcase the potential of AI in a variety of scenarios. However, as was discussed earlier in the session, centralised training is unsuitable for a larger range of settings were information cannot be collected in order to build a unified (centralised) dataset.\n",
    "\n",
    "# The Future of AI is Federated\n",
    "\n",
    "What are the key differences between Federated Learning and Centralised Training? I you could only pick you, probably you'd say:\n",
    "* Federated Learning is distributed -- the model is trained on-device by the participating clients.\n",
    "* Data remains private and is owned by a specific _client_ -- the data is never sent to the central server.\n",
    "\n",
    "The are several more differences. But the above two are the main ones to always consider and that are common to all flavours of Federated Learning (e.g. _cross-device_ or _cross-silo_). The remaining of this tutorial is going to focus in transforming the code we have written so far for the centralised setting and construct a Federated Learning pipeline using Flower and PyTorch.\n",
    "\n",
    "Let's begin! 🚀"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One Client, One Data Partition\n",
    "\n",
    "To start designing a Federated Learning pipeline we need to meet one of the key properties in FL: each client has its own data partition. To accomplish this with the MNIST dataset, we are going to generate N random partitions, where N is the total number of clients in our FL system, using [Flower Datasets](https://flower.ai/docs/datasets/). Let's create 100 partitions with the [IidPartitioner](https://flower.ai/docs/datasets/ref-api/flwr_datasets.partitioner.IidPartitioner.html#flwr_datasets.partitioner.IidPartitioner) -- note there are many more [partitioners](https://flower.ai/docs/datasets/ref-api/flwr_datasets.partitioner.html) to choose from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr_datasets import FederatedDataset\n",
    "from flwr_datasets.partitioner import IidPartitioner\n",
    "\n",
    "NUM_PARTITIONS = 100\n",
    "\n",
    "partitioner = IidPartitioner(num_partitions=NUM_PARTITIONS)\n",
    "# Let's partition the \"train\" split of the MNIST dataset\n",
    "# The MNIST dataset will be downloaded if it hasn't been already\n",
    "fds = FederatedDataset(dataset=\"ylecun/mnist\", partitioners={\"train\": partitioner})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Accessing individual partitions can be done like this. The return object can be then passed to a dataloader for training or evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We could load a single partition like this\n",
    "partition_0 = fds.load_partition(0)\n",
    "partition_0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Flower Datasets` comes with built-in visualization tools that help you get insights of how the dataset (in this case MNIST) has been partitioned. Let's create a parplot to visualize the number of labels of each class that every client's partition contains. Note we are only visualising the first 30 clients purely so the plot remain readable. \n",
    "\n",
    "> There are many more types of plots you can generated with Flower Datasets. Check the [Visualization tutorial](https://flower.ai/docs/datasets/tutorial-visualize-label-distribution.html). Feel free to try other partitioning scheemes and you'll see how the visualization changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr_datasets.visualization import plot_label_distributions\n",
    "\n",
    "fig, ax, df = plot_label_distributions(\n",
    "    partitioner,\n",
    "    label_name=\"label\",\n",
    "    plot_type=\"bar\",\n",
    "    size_unit=\"absolute\",\n",
    "    partition_id_axis=\"x\",\n",
    "    legend=True,\n",
    "    verbose_labels=True,\n",
    "    max_num_partitions=30,  # Note we are only showing the first 30 so the plot remains readable\n",
    "    title=\"Per Partition Labels Distribution\",\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the histogram of this partition is a bit different from the one we obtained at the beginning where we took the entire dataset into consideration. Because our data partitions are artificially constructed by sampling the MNIST dataset in an IID fashion, our Federated Learning example will not face sever _data heterogeneity_ issues (which is a fairly [active research topic](https://arxiv.org/abs/1912.04977)).\n",
    "\n",
    "Let's next define how our FL clients will behave\n",
    "\n",
    "## Defining a Flower `ClientApp`\n",
    "\n",
    "You can think of a client in FL as an entity that owns some data and trains a model using this data. The caveat is that the model is being trained _collaboratively_ in Federation by multiple clients (sometimes up to hundreds of thousands) and, in most instances of FL, is sent by a central server running in a `ServerApp` (more on this later).\n",
    "\n",
    "A Flower Client is a simple Python class with two distinct methods:\n",
    "\n",
    "* `fit()`: With this method, the client does on-device training for a number of epochs using its own data. At the end, the resulting model is sent back to the server for aggregation.\n",
    "\n",
    "* `evaluate()`: With this method, the server can evaluate the performance of the global model on the local validation set of a client. This can be used for instance when there is no centralised dataset on the server for validation/test. Also, this method can be use to asses the degree of personalisation of the model being federated.\n",
    "\n",
    "This class will be then wrapped into a `ClientApp` that can be used to launch the simulation."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define our Flower Client class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from typing import Dict, Tuple\n",
    "\n",
    "import torch\n",
    "from flwr.common import NDArrays, Scalar\n",
    "from flwr.client import NumPyClient\n",
    "\n",
    "\n",
    "class FlowerClient(NumPyClient):\n",
    "    def __init__(self, trainloader, valloader) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.trainloader = trainloader\n",
    "        self.valloader = valloader\n",
    "        self.model = Net(num_classes=10)\n",
    "\n",
    "    def fit(self, parameters, config):\n",
    "        \"\"\"This method trains the model using the parameters sent by the\n",
    "        server on the dataset of this client. At then end, the parameters\n",
    "        of the locally trained model are communicated back to the server\"\"\"\n",
    "\n",
    "        # copy parameters sent by the server into client's local model\n",
    "        set_params(self.model, parameters)\n",
    "\n",
    "        # Define the optimizer\n",
    "        optim = torch.optim.SGD(self.model.parameters(), lr=0.01, momentum=0.9)\n",
    "\n",
    "        # do local training (call same function as centralised setting)\n",
    "        train(self.model, self.trainloader, optim, epochs=1)\n",
    "\n",
    "        # return the model parameters to the server as well as extra info (number of training examples in this case)\n",
    "        return get_params(self.model), len(self.trainloader), {}\n",
    "\n",
    "    def evaluate(self, parameters: NDArrays, config: Dict[str, Scalar]):\n",
    "        \"\"\"Evaluate the model sent by the server on this client's\n",
    "        local validation set. Then return performance metrics.\"\"\"\n",
    "\n",
    "        set_params(self.model, parameters)\n",
    "        # do local evaluation (call same function as centralised setting)\n",
    "        loss, accuracy = test(self.model, self.valloader)\n",
    "        # send statistics back to the server\n",
    "        return float(loss), len(self.valloader), {\"accuracy\": accuracy}\n",
    "\n",
    "\n",
    "# Two auxhiliary functions to set and extract parameters of a model\n",
    "def set_params(model, parameters):\n",
    "    \"\"\"Replace model parameters with those passed as `parameters`.\"\"\"\n",
    "\n",
    "    params_dict = zip(model.state_dict().keys(), parameters)\n",
    "    state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})\n",
    "    # now replace the parameters\n",
    "    model.load_state_dict(state_dict, strict=True)\n",
    "\n",
    "\n",
    "def get_params(model):\n",
    "    \"\"\"Extract model parameters as a list of NumPy arrays.\"\"\"\n",
    "    return [val.cpu().numpy() for _, val in model.state_dict().items()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Spend a few minutes to inspect the `FlowerClient` class above. Please ask questions if there is something unclear !\n",
    "\n",
    "Then keen-eyed among you might have realised that if we were to fuse the client's `fit()` and `evaluate()` methods, we'll end up with essentially the same as in the `run_centralised()` function we used in the Centralised Training part of this tutorial. And it is true!! In Federated Learning, the way clients perform local training makes use of the same principles as more traditional centralised setup. The key difference is that the dataset now is much smaller and it's never _\"seen\"_ by the entity running the FL workload (i.e. the central server).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The `client_fn` callback\n",
    "\n",
    "Now let's see how the `FlowerClient` object above can be used in Flower: we need to construct a `ClientApp`. This can be conveniently be done by means of a `client_fn` callback that will return a `FlowerClient` that uses a specific data partition (`partition-id`). The index of the partition is set internally during the simulation (meaning you shouldn't worry about it this tutorial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.common import Context\n",
    "from flwr.client import ClientApp\n",
    "\n",
    "\n",
    "def client_fn(context: Context):\n",
    "    \"\"\"Returns a FlowerClient containing its data partition.\"\"\"\n",
    "\n",
    "    partition_id = int(context.node_config[\"partition-id\"])\n",
    "    partition = fds.load_partition(partition_id, \"train\")\n",
    "    # partition into train/validation\n",
    "    partition_train_val = partition.train_test_split(test_size=0.1, seed=42)\n",
    "\n",
    "    # Let's use the function defined earlier to construct the dataloaders\n",
    "    # and apply the dataset transformations\n",
    "    trainloader, testloader = get_mnist_dataloaders(partition_train_val, batch_size=32)\n",
    "\n",
    "    return FlowerClient(trainloader=trainloader, valloader=testloader).to_client()\n",
    "\n",
    "\n",
    "# Concstruct the ClientApp passing the client generation function\n",
    "client_app = ClientApp(client_fn=client_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that a `ClientApp` is fully defined, let's create its counterpart: the `ServerApp`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining a Flower `ServerApp`\n",
    "\n",
    "### Chosing a Flower Strategy\n",
    "\n",
    "\n",
    "A strategy sits at the core of the Federated Learning experiment. It is involved in all stages of a FL pipeline: sampling clients; sending the _global model_ to the clients so they can do `fit()`; receive the updated models from the clients and **aggregate** these to construct a new _global model_; define and execute global or federated evaluation; and more.\n",
    "\n",
    "Flower comes with [many strategies built-in](https://github.com/adap/flower/tree/main/src/py/flwr/server/strategy). For this tutorial, let's use what is arguable the most popular strategy out there: `FedAvg`.\n",
    "\n",
    "The way `FedAvg` works is simple but performs surprisingly well in practice. It is therefore one good strategy to start your experimentation. `FedAvg`, as its name implies, derives a new version of the _global model_ by taking the average of all the models sent by clients participating in the round. You can read all the details [in the paper](https://arxiv.org/abs/1602.05629).\n",
    "\n",
    "While Flower strategies offer a high degree of customization using callbacks, in this tutorial we'll focus on using just one: the `evaluate_metrics_aggregation_fn` callback. It allows you to pass a function that should be executed at the end of an _\"evaluate\"_ round (i.e. a round where clients evaluate the _global model_ they receive on their local data and report the result -- e.g. accuracy, loss, etc -- back to the server). For this tutorial we want to perform the weighted average of the _\"accuracy\"_ metrics returned by each `FlowerClient`'s `evaluate()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from flwr.common import Metrics\n",
    "\n",
    "\n",
    "# Define metric aggregation function\n",
    "def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:\n",
    "    # Multiply accuracy of each client by number of examples used\n",
    "    accuracies = [num_examples * m[\"accuracy\"] for num_examples, m in metrics]\n",
    "    examples = [num_examples for num_examples, _ in metrics]\n",
    "\n",
    "    # Aggregate and return custom metric (weighted average)\n",
    "    return {\"accuracy\": sum(accuracies) / sum(examples)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use this callback when defining the strategy in the next section"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The `server_fn` callback\n",
    "\n",
    "The easiest way to create a `ServerApp` with the aggregation _strategy_ of your choice is by means of a `server_fn` callback. It has a similar signature to `client_fn` but, instead of returning a client object, it returns all the components needed to run the server-side logic in Flower. In this tutorial we'll keep things simple and stick to `FedAvg` with initialised global parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.common import ndarrays_to_parameters\n",
    "from flwr.server import ServerApp, ServerConfig, ServerAppComponents\n",
    "from flwr.server.strategy import FedAvg\n",
    "\n",
    "num_rounds = 5\n",
    "\n",
    "\n",
    "def server_fn(context: Context):\n",
    "\n",
    "    # instantiate the model\n",
    "    model = Net(num_classes=10)\n",
    "    ndarrays = get_params(model)\n",
    "    # Convert model parameters to flwr.common.Parameters\n",
    "    global_model_init = ndarrays_to_parameters(ndarrays)\n",
    "\n",
    "    # Define the strategy\n",
    "    strategy = FedAvg(\n",
    "        fraction_fit=0.1,  # 10% clients sampled each round to do fit()\n",
    "        fraction_evaluate=0.5,  # 50% clients sample each round to do evaluate()\n",
    "        evaluate_metrics_aggregation_fn=weighted_average,  # callback defined earlier\n",
    "        initial_parameters=global_model_init,  # initialised global model\n",
    "    )\n",
    "\n",
    "    # Construct ServerConfig\n",
    "    config = ServerConfig(num_rounds=num_rounds)\n",
    "\n",
    "    # Wrap everything into a `ServerAppComponents` object\n",
    "    return ServerAppComponents(strategy=strategy, config=config)\n",
    "\n",
    "\n",
    "# Create your ServerApp\n",
    "server_app = ServerApp(server_fn=server_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launching the Simulation\n",
    "\n",
    "With both `ClientApp` and `ServerApp` ready, we can launch the simulation. Pass both apps to the `run_simulation()` function and specify the number of `supernodes` (this is a more general term used in Flower to refer to individual \"nodes\" or \"clients\"). We earlier partitioned the dataset into 100 partitions, one for each supernode. So we indicate that `num_supernodes`=100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.simulation import run_simulation\n",
    "\n",
    "run_simulation(\n",
    "    server_app=server_app, client_app=client_app, num_supernodes=NUM_PARTITIONS\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note above how the distributed `accuracy` goes up as training progresses while the loss goes down. Federated learning is working!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus: Create your own Strategy\n",
    "\n",
    "Flower strategies can be extended easily to suit your FL setups or your preferred workflows whether you use Flower for research or in production. In this final section, you'll learn how to create a custom strategy that behaves just like `FedAvg` but extends the functionality of certain methods to achieve two things:\n",
    "1. Save the results obtained on each round into a JSON file.\n",
    "2. Create a plot at after the last round.\n",
    "\n",
    "\n",
    "Let's call this strategy `FedAvgCustom`. We'll use it to also showcase how to use the `evaluate_fn` callback, a convenient way to do centralised evaluation of the global model after each round. Note this functionality is user-defined since it requires a choice in terms of ML-framework. (if you recall, Flower is framework agnostic).\n",
    "\n",
    "> This being said, centralised evaluation of the global model is only possible if there exists a centralised dataset that somewhat follows a similar distribution as the data that's spread across clients. In some cases having such centralised dataset for validation is not possible, so the only solution is to federate the evaluation of the _global model_. This is the default behaviour in Flower. If you don't specify the `evaluate_fn` argument in your strategy, then, centralised global evaluation won't be performed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.server.strategy import FedAvg\n",
    "from flwr.common import Parameters\n",
    "import json\n",
    "\n",
    "\n",
    "class FedAvgCustom(FedAvg):\n",
    "    def __init__(self, file_name: str, num_rounds: int, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.file_name = file_name\n",
    "        self.num_rounds = num_rounds\n",
    "        self.loss_list = []\n",
    "        self.metrics_list = []\n",
    "\n",
    "    def _make_plot(self):\n",
    "        \"\"\"Makes a plot with the results recorded\"\"\"\n",
    "        round = list(range(1, len(self.loss_list) + 1))\n",
    "        acc = [100.0 * metrics[\"accuracy\"] for metrics in self.metrics_list]\n",
    "        plt.plot(round, acc)\n",
    "        plt.grid()\n",
    "        plt.ylabel(\"Accuracy (%)\")\n",
    "        plt.xlabel(\"Round\")\n",
    "\n",
    "    def evaluate(self, server_round: int, parameters: Parameters):\n",
    "        \"\"\"Evaluate model parameters using an evaluation function.\"\"\"\n",
    "        loss, metrics = super().evaluate(server_round, parameters)\n",
    "        # Record results\n",
    "        self.loss_list.append(loss)\n",
    "        self.metrics_list.append(metrics)\n",
    "        # If last round, save results and make a plot\n",
    "        if server_round == self.num_rounds:\n",
    "            # Save to CSV\n",
    "            with open(f\"{self.file_name}.json\", \"w\") as f:\n",
    "                json.dump({\"loss\": self.loss_list, \"metrics\": self.metrics_list}, f)\n",
    "            # Generate plot\n",
    "            self._make_plot()\n",
    "\n",
    "\n",
    "def get_evaluate_fn(testloader):\n",
    "    \"\"\"Return a function that can be called to do global evaluation.\"\"\"\n",
    "\n",
    "    def evaluate_fn(server_round: int, parameters, config):\n",
    "        \"\"\"Evaluate global model on the whole test set.\"\"\"\n",
    "\n",
    "        model = Net(num_classes=10)\n",
    "\n",
    "        # set parameters to the model\n",
    "        params_dict = zip(model.state_dict().keys(), parameters)\n",
    "        state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})\n",
    "        model.load_state_dict(state_dict, strict=True)\n",
    "\n",
    "        # call test (evaluate model as in centralised setting)\n",
    "        loss, accuracy = test(model, testloader)\n",
    "        return loss, {\"accuracy\": accuracy}\n",
    "\n",
    "    return evaluate_fn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the above defined, we just need to wrap it all up in a `ServerApp` as we did earlier but this time using the `FedAvgCustom` that we just defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "9ad8dcea-8004-4c6e-a025-e168da636c88"
   },
   "outputs": [],
   "source": [
    "from flwr.server import ServerApp, ServerConfig\n",
    "\n",
    "\n",
    "def server_fn(context: Context):\n",
    "\n",
    "    # instantiate the model\n",
    "    model = Net(num_classes=10)\n",
    "    ndarrays = get_params(model)\n",
    "    # Convert model parameters to flwr.common.Parameters\n",
    "    global_model_init = ndarrays_to_parameters(ndarrays)\n",
    "\n",
    "    # Define the strategy\n",
    "    strategy = FedAvgCustom(\n",
    "        file_name=\"results_fedavgcustom\",\n",
    "        num_rounds=num_rounds,\n",
    "        fraction_fit=0.1,  # 10% clients sampled each round to do fit()\n",
    "        fraction_evaluate=0.25,  # 25% clients sample each round to do evaluate()\n",
    "        evaluate_metrics_aggregation_fn=weighted_average,  # callback defined earlier\n",
    "        initial_parameters=global_model_init,  # initialised global model\n",
    "        evaluate_fn=get_evaluate_fn(\n",
    "            testloader\n",
    "        ),  # gloabl evaluation (here we can pass the same testset as used in centralised)\n",
    "    )\n",
    "\n",
    "    # Construct ServerConfig\n",
    "    config = ServerConfig(num_rounds=num_rounds)\n",
    "\n",
    "    # Wrap everything into a `ServerAppComponents` object\n",
    "    return ServerAppComponents(strategy=strategy, config=config)\n",
    "\n",
    "\n",
    "# Create your ServerApp\n",
    "server_app = ServerApp(server_fn=server_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All that is left is to launch the simulation. Note a plot will be displayed at the end and a `.json` with the results will be saved to the current directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_simulation(\n",
    "    server_app=server_app, client_app=client_app, num_supernodes=NUM_PARTITIONS\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's next?\n",
    "\n",
    "Well, if you enjoyed this content, consider giving us a ⭐️ on GitHub -> https://github.com/adap/flower\n",
    "\n",
    "* **[DOCS]** How about running your Flower clients on the GPU? find out how to do it in the [Flower Simulation Documentation](https://flower.ai/docs/framework/how-to-run-simulations.html)\n",
    "\n",
    "* **[VIDEO]** You can follow our [detailed line-by-line 9-videos tutorial](https://www.youtube.com/watch?v=cRebUIGB5RU&list=PLNG4feLHqCWlnj8a_E1A_n5zr2-8pafTB) about everything you need to know to design your own Flower Simulation pipelines\n",
    "\n",
    "* Check more advanced simulation examples the Flower GitHub:\n",
    "\n",
    "    * Flower simulation with Tensorflow/Keras: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/adap/flower/tree/main/examples/simulation-tensorflow)\n",
    "\n",
    "    * Flower simulation with Pytorch: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/adap/flower/tree/main/examples/simulation-pytorch)\n",
    "\n",
    "* **[DOCS]** All Flower examples: https://flower.ai/docs/examples/\n",
    "\n",
    "* **[VIDEO]** Our Youtube channel: https://www.youtube.com/@flowerlabs\n",
    "\n",
    "Don't forget to join our Slack channel: https://flower.ai/join-slack/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
